from filestructs import StructReader, TreeNode, Forest, ListFile
from codepoints import CodePoint, FunctionCP, BuiltinCP, LogCP, RootCP
from struct import Struct


UI16 = Struct("=H")
FUNC_ENTER = Struct("=QHQ")
FUNC_LEAVE = Struct("=QQ")


class TraceForest(object):
    @classmethod
    def load(cls, fileprefix):
        root = Forest.open(fileprefix + ".trace")
        codepoints = CodePoint.from_filename(fileprefix + ".cp")
        values = ListFile.open(fileprefix + ".values")
        return cls.create(root, codepoints, values)

    @classmethod
    def create(cls, root, codepoints, values):
        subclasses = {}
        class TraceNode(object):
            __slots__ = ["node", "cp", "_children"]
            REPR_ATTRS = ["cp"]
            
            def __init__(self, node, cp):
                self.node = node
                self.cp = cp
                self._children = None
            
            @classmethod
            def create(cls, node):
                cpindex, = UI16.unpack(node.data[:2])
                cp = codepoints[cpindex]
                return subclasses[type(cp)](node, cp)
            
            def _preload(self):
                if self._children is not None:
                    return
                self._children = [self.create(child) for child in self.node]
            
            def __len__(self):
                self._preload()
                return len(self._children)
            def __iter__(self):
                self._preload()
                return iter(self._children)
            def __getitem__(self, index):
                self._preload()
                return self._children[index]
            def has_children(self):
                return self.node.has_children
            
            def __repr__(self):
                attrs = [repr(getattr(self, name)) for name in self.REPR_ATTRS]
                return "%s(%s)" % (self.__class__.__name__, ", ".join(attrs))
            
            def unload(self):
                if self._children is None:
                    return
                for child in self._children:
                    child.unload()
                self._children = None
            
            def show(self, level = 0):
                print "%s%s" % ("    " * level, self)
                level += 1
                for child in self:
                    child.show(level)
        
        class LogNode(TraceNode):
            #__slots__ = ["payload"]
            def __init__(self, node, cp):
                TraceNode.__init__(self, node, cp)
                self.payload = node.data[2:]
                self.fucntion = "--LOG--"
                self.lineno = 0
                self.filename = "/dev/null"
        
        class EntryInfo(object):
            __slots__ = ["timestamp", "argcount", "argoffset"]
            def __init__(self, blob):
                self.timestamp, self.argcount, self.argoffset = FUNC_ENTER.unpack(blob)
        
        class LeaveInfo(object):
            __slots__ = ["timestamp", "retval", "rettype"]
            def __init__(self, blob):
                self.timestamp, retval  = FUNC_LEAVE.unpack(blob)
                self.rettype = retval & 0x0f # 4 lower bits
                self.retval = retval >> 4 # 60 upper bits
        
        class FuncNode(TraceNode):
            #__slots__ = ["entry", "leave"]
            def __init__(self, node, cp):
                TraceNode.__init__(self, node, cp)
                blob1 = node.data[2:2+FUNC_ENTER.size]
                blob2 = node.data[2+FUNC_ENTER.size:]
                self.entry = EntryInfo(blob1)
                if blob2:
                    self.leave = LeaveInfo(blob2)
                else:
                    self.leave = None
                self.function = cp.name
                self.lineno = getattr(cp, "lineno", -1)
                self.filename = getattr(cp, "filename", "?")
        
        subclasses[LogCP] = LogNode
        subclasses[FunctionCP] = FuncNode
        subclasses[BuiltinCP] = FuncNode
        
        class _TraceForest(object):
            __slots__ = ["root", "_children"]
            def __init__(self, root):
                self.root = root
                self._children = None
            def _preload(self):
                if self._children is not None:
                    return
                self._children = [TraceNode.create(node) for node in self.root]
            
            def has_children(self):
                self._preload()
                return bool(self._children)
            def __len__(self):
                self._preload()
                return len(self._children)
            def __iter__(self):
                self._preload()
                return iter(self._children)
            def __getitem__(self, index):
                self._preload()
                return self._children[index]
            
            def show(self):
                print "TraceForest"
                for node in self:
                    node.show(1)
        
        return _TraceForest(root)



if __name__ == "__main__":
    #fn = "/a/home/tomer/xiv/tlib/runner.py-20090129125552-0"
    fn = "/a/home/tomer/xiv/tlib/runner.py-20090129135232-0"
    #fn = "../../build/lib.linux-i686-2.5/test2.py-20090127192240"
    root = TraceForest.load(fn)
    root.show()











